import tensorflow as tf
import math
from tensorflow.examples.tutorials.mnist import input_data as mnist_data
print("Tensorflow version " + tf.__version__)
tf.set_random_seed(0)
import numpy as np
import cv2

mnist = mnist_data.read_data_sets("data", one_hot=True, reshape=False, validation_size=0)
X = tf.placeholder(tf.float32, [None, 28, 28, 1])
Y_= tf.placeholder(tf.float32, [None, 10])
# variable learning rate
lr= tf.placeholder(tf.float32)
# Probability of keeping a node during dropout = 1.0 at test time (no dropout) and 0.75 at training time
pkeep = tf.placeholder(tf.float32)
# five layers and their number of neurons (tha last layer has 10 softmax neurons)
L = 200
M = 100
N = 60
O = 30
# Weights initialised with small random values between -0.2 and +0.2
# When using RELUs, make sure biases are initialised with small *positive* values for example 0.1 = tf.ones([K])/10
W1 = tf.Variable(tf.truncated_normal([784, L], stddev=0.1))  # 784 = 28 * 28
B1 = tf.Variable(tf.ones([L])/10)
W2 = tf.Variable(tf.truncated_normal([L, M], stddev=0.1))
B2 = tf.Variable(tf.ones([M])/10)
W3 = tf.Variable(tf.truncated_normal([M, N], stddev=0.1))
B3 = tf.Variable(tf.ones([N])/10)
W4 = tf.Variable(tf.truncated_normal([N, O], stddev=0.1))
B4 = tf.Variable(tf.ones([O])/10)
W5 = tf.Variable(tf.truncated_normal([O, 10], stddev=0.1))
B5 = tf.Variable(tf.zeros([10]))
# The model, with dropout at each layer
XX = tf.reshape(X, [-1, 28*28])
Y1 = tf.nn.relu(tf.matmul(XX, W1) + B1)
Y1d = tf.nn.dropout(Y1, pkeep)
Y2 = tf.nn.relu(tf.matmul(Y1d, W2) + B2)
Y2d = tf.nn.dropout(Y2, pkeep)
Y3 = tf.nn.relu(tf.matmul(Y2d, W3) + B3)
Y3d = tf.nn.dropout(Y3, pkeep)
Y4 = tf.nn.relu(tf.matmul(Y3d, W4) + B4)
Y4d = tf.nn.dropout(Y4, pkeep)
Ylogits = tf.matmul(Y4d, W5) + B5
Y = tf.nn.softmax(Ylogits)
# cross-entropy loss function (= -sum(Y_i * log(Yi)) ), normalised for batches of 100  images
# TensorFlow provides the softmax_cross_entropy_with_logits function to avoid numerical stability
# problems with log(0) which is NaN
cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits=Ylogits, labels=Y_)
cross_entropy = tf.reduce_mean(cross_entropy)*100
# accuracy of the trained model, between 0 (worst) and 1 (best)
correct_prediction = tf.equal(tf.argmax(Y, 1), tf.argmax(Y_, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
# training step, the learning rate is a placeholder
train_step = tf.train.AdamOptimizer(lr).minimize(cross_entropy)
# init
init = tf.global_variables_initializer()
# save model
Saver = tf.train.Saver(max_to_keep = 1)  # defaults to saving all variables
ModelSavePath = './storage/'
def train():
	with tf.Session() as sess:
		sess.run(init)
		ckpt = tf.train.get_checkpoint_state(ModelSavePath)
		if ckpt and ckpt.model_checkpoint_path:
			Saver.restore(sess, ckpt.model_checkpoint_path)
			print '----> restore model ok'
		else:
			print '----> no model so far'
		iterNum = 50000
		for i in xrange(iterNum):
			if (i + 1) % 100 == 0:
				Saver.save(sess, ModelSavePath + "model", global_step = (i + 1))
			# training on batches of 100 images with 100 labels
			batch_X, batch_Y = mnist.train.next_batch(100)
			# learning rate decay
			max_learning_rate = 0.003
			min_learning_rate = 0.0001
			decay_speed = 2000.0 # 0.003-0.0001-2000=>0.9826 done in 5000 iterations
			learning_rate = min_learning_rate + (max_learning_rate - min_learning_rate) * math.exp(-i/decay_speed)
			a, c, s = sess.run([accuracy, cross_entropy, train_step], \
			{X: batch_X, Y_: batch_Y, pkeep: 0.75, lr: learning_rate})
			if (i + 1) % 100 == 0:
				print iterNum, 	'%6d'%i, \
							'%9.3f(cross entropy) '%c, \
							'%9.3f(training accuracy) '%a, \
							'%9.5f(learning rate) '%learning_rate

def test():
	batch_X = np.zeros([6, 28, 28, 1])
	batch_X[0] = np.expand_dims(cv2.imread('../z0.jpg', cv2.IMREAD_GRAYSCALE).astype('float32')/ 255.0, -1)
	batch_X[1] = np.expand_dims(cv2.imread('../z1.jpg', cv2.IMREAD_GRAYSCALE).astype('float32')/ 255.0, -1)
	batch_X[2] = np.expand_dims(cv2.imread('../z2.jpg', cv2.IMREAD_GRAYSCALE).astype('float32')/ 255.0, -1)
	batch_X[3] = np.expand_dims(cv2.imread('../z3.jpg', cv2.IMREAD_GRAYSCALE).astype('float32')/ 255.0, -1)
	batch_X[4] = np.expand_dims(cv2.imread('../z4.jpg', cv2.IMREAD_GRAYSCALE).astype('float32')/ 255.0, -1)
	batch_X[5] = np.expand_dims(cv2.imread('../z5.jpg', cv2.IMREAD_GRAYSCALE).astype('float32')/ 255.0, -1)
	trueY = np.array([6, 7, 2, 4, 3, 5])

	with tf.Session() as sess:
		sess.run(init)
		ckpt = tf.train.get_checkpoint_state(ModelSavePath)
		if ckpt and ckpt.model_checkpoint_path:
			Saver.restore(sess, ckpt.model_checkpoint_path)
			print '----> restore model ok'
		else:
			print '----> no model so far'
		oy = sess.run(Y, {X: batch_X, pkeep: 1.0})
		for i in range(0, 6):
			for j in range(0, 10):
				if(oy[i, j] > 0.6):
					print i, trueY[i], j
					break

def test1():
	batch_X, batch_Y = mnist.train.next_batch(1)
	img = batch_X[0,:,:,0] * 255.0
	img = img.astype('uint8')
	print img
#	cv2.imshow('win', img)
#	cv2.waitKey(0)
	print batch_Y
	
	with tf.Session() as sess:
		sess.run(init)
		ckpt = tf.train.get_checkpoint_state(ModelSavePath)
		if ckpt and ckpt.model_checkpoint_path:
			Saver.restore(sess, ckpt.model_checkpoint_path)
			print '----> restore model ok'
		else:
			print '----> no model so far'
		oy = sess.run(Y, {X: batch_X, pkeep: 1.0})
		print oy

if __name__ == "__main__":
#	train()
	test()

